"""Alibaba Qwen (通义千问) model adapter implementation."""

from typing import AsyncGenerator, Dict, Any, List
import json
import aiohttp
from dashscope import Generation
import dashscope

from .base import BaseModelAdapter, ModelConfig, ModelResponse


class QwenAdapter(BaseModelAdapter):
    """Adapter for Alibaba Qwen models."""
    
    # Pricing per 1K tokens (as of 2024)
    PRICING = {
        "qwen-max": {"prompt": 0.02, "completion": 0.06},
        "qwen-plus": {"prompt": 0.004, "completion": 0.012},
        "qwen-turbo": {"prompt": 0.002, "completion": 0.006}
    }
    
    def _initialize_client(self) -> None:
        """Initialize DashScope client."""
        dashscope.api_key = self.config.api_key
        if self.config.api_base:
            dashscope.base_url = self.config.api_base
    
    async def generate(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> ModelResponse:
        """Generate response from Qwen model.
        
        Args:
            messages: Conversation messages
            **kwargs: Additional parameters
            
        Returns:
            ModelResponse with generated content
        """
        params = {
            "model": self.config.model_name,
            "messages": messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "top_p": self.config.top_p,
            "result_format": "message"
        }
        
        params.update(kwargs)
        params.update(self.config.extra_params)
        
        # DashScope SDK is synchronous, run in executor
        import asyncio
        response = await asyncio.get_event_loop().run_in_executor(
            None,
            Generation.call,
            **params
        )
        
        if response.status_code != 200:
            raise Exception(f"Qwen API error: {response.message}")
        
        # Parse response
        content = response.output.choices[0].message.content
        usage = {
            "prompt_tokens": response.usage.input_tokens,
            "completion_tokens": response.usage.output_tokens,
            "total_tokens": response.usage.total_tokens
        }
        
        return ModelResponse(
            content=content,
            model=self.config.model_name,
            usage=usage,
            finish_reason=response.output.choices[0].finish_reason,
            metadata={
                "request_id": response.request_id
            }
        )
    
    async def stream(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> AsyncGenerator[str, None]:
        """Stream response from Qwen model.
        
        Args:
            messages: Conversation messages
            **kwargs: Additional parameters
            
        Yields:
            Content chunks
        """
        # Prepare request
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json",
            "X-DashScope-SSE": "enable"
        }
        
        api_base = self.config.api_base or "https://dashscope.aliyuncs.com/api/v1"
        url = f"{api_base}/services/aigc/text-generation/generation"
        
        data = {
            "model": self.config.model_name,
            "input": {"messages": messages},
            "parameters": {
                "temperature": self.config.temperature,
                "max_tokens": self.config.max_tokens,
                "top_p": self.config.top_p,
                "incremental_output": True,
                "result_format": "message"
            }
        }
        
        data["parameters"].update(kwargs)
        data["parameters"].update(self.config.extra_params)
        
        # Stream response
        async with aiohttp.ClientSession() as session:
            async with session.post(
                url,
                headers=headers,
                json=data,
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            ) as response:
                async for line in response.content:
                    if line:
                        line_str = line.decode('utf-8').strip()
                        if line_str.startswith("data:"):
                            data_str = line_str[5:].strip()
                            if data_str and data_str != "[DONE]":
                                try:
                                    chunk_data = json.loads(data_str)
                                    if "output" in chunk_data and "choices" in chunk_data["output"]:
                                        choice = chunk_data["output"]["choices"][0]
                                        if "message" in choice and "content" in choice["message"]:
                                            yield choice["message"]["content"]
                                except json.JSONDecodeError:
                                    continue
    
    def count_tokens(self, text: str) -> int:
        """Count tokens for Qwen models.
        
        Rough estimate: 1 Chinese char ≈ 1 token, 1 English word ≈ 1.5 tokens
        
        Args:
            text: Input text
            
        Returns:
            Estimated token count
        """
        # Count Chinese characters
        chinese_chars = sum(1 for char in text if '\u4e00' <= char <= '\u9fff')
        # Count English words (rough estimate)
        english_words = len([w for w in text.split() if not any('\u4e00' <= c <= '\u9fff' for c in w)])
        
        return chinese_chars + int(english_words * 1.5)
    
    def estimate_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
        """Estimate cost for Qwen models.
        
        Args:
            prompt_tokens: Number of prompt tokens
            completion_tokens: Number of completion tokens
            
        Returns:
            Cost in USD (converted from CNY)
        """
        model_key = self.config.model_name
        
        # Find matching pricing
        for key in self.PRICING:
            if key in model_key:
                model_key = key
                break
        else:
            # Default to qwen-turbo if unknown
            model_key = "qwen-turbo"
        
        pricing = self.PRICING[model_key]
        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
        completion_cost = (completion_tokens / 1000) * pricing["completion"]
        
        # Prices are in CNY, convert to USD (rough rate)
        cny_to_usd = 0.14
        return (prompt_cost + completion_cost) * cny_to_usd